Description

Discussed with @FBruzzesi a bit on Discord, so I thought I'd try opening a PR to show what this would look like. Just starting off with the pandas_utils function, curious to get some initial feedback and to know if there'd be interest in this before moving on to the rest

The idea is that, instead of hard-coding support for pandas, you could make use of the DataFrame Consortium API, which is (wip in progress!) defined here: https://data-apis.org/dataframe-api/draft/index.html

Then, any function which adheres to that API will:

• keep support for pandas (so, no regression in functionality)
• also gain support for Polars, and any other dataframe library which implements the API (devs from cuDF and Modin also want to implement it)

Example

If you run

import pandas as pd
import polars as pl

from sklego.pandas_utils import add_lags

df = pd.DataFrame({"X1": [0, 1, 2], "X2": [float('nan'), "178", "154"]})
print(add_lags(df, "X1", -1))

then it will print

   X1   X2  X1-1
0   1  178   0.0
1   2  154   1.0

, just like it does now. No change here.

However, you can now also run

import polars as pl

from sklego.pandas_utils import add_lags

df = pl.DataFrame({"X1": [0, 1, 2], "X2": [float('nan'), "178", "154"]})
result = add_lags(df, "X1", -1)
print(result.collect())

and you'll get

shape: (2, 3)
┌─────┬─────┬──────┐
│ X1  ┆ X2  ┆ X1-1 │
│ --- ┆ --- ┆ ---  │
│ i64 ┆ str ┆ i64  │
╞═════╪═════╪══════╡
│ 1   ┆ 178 ┆ 0    │
│ 2   ┆ 154 ┆ 1    │
└─────┴─────┴──────┘

(note how I had to add collect before printing the result, as the result is a Polars LazyFrame)

Dependencies

Note how in sklego/dataframe_utils.py, it's now possible to remove import pandas as pd. If this API were used throughout the package, then you wouldn't even need pandas as a required dependency - scikit-lego would be truly dataframe-agnostic, and would use whichever dataframe package the user passes.

All that would be needed would the dataframe-api-compat package. Note that it's light as a feather:

• it's pure Python, with no dependencies (not even pandas / Polars, it'll just defer to whatever the user calls __dataframe_consortium_standard__ on)
• the wheel is literally only about 20-30 kilobytes https://pypi.org/project/dataframe-api-compat/0.1.28/#files
• it would be easy to vendor if you didn't want to take on an extra dependency

Checklist:

• My code follows the style guidelines (flake8)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (also to the readme.md)
• I have added tests that prove my fix is effective or that my feature works
• I have added tests to check whether the new feature adheres to the sklearn convention
• New and existing unit tests pass locally with my changes

If you feel your PR is ready for a review, ping @koaning or @MBrouns.